#!/usr/bin/env python3
# coding=utf-8
'''
code is based on DUI official doc at:
    https://www.duiopen.com/docs/ct_asr_sentence
'''

import os, sys
import argparse
import time
import subprocess

import json


def arr_size(arr, size):
    s = []
    for i in range(0, len(arr) + 1, size):
        c = arr[i:i + size]
        s.append(c)
    newlist = [x for x in s if x]
    return newlist


sherpa = "sherpa-onnx-v1.10.38-linux-x64-static"
sherpa_dir = os.getcwd() + f"/{sherpa}"
model = "sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20"

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('scp_path', type=str)
    parser.add_argument('trans_path', type=str)
    args = parser.parse_args()
    cmd = f"{sherpa_dir}/bin/sherpa-onnx --num-threads=16 --tokens=./{sherpa}/{model}/tokens.txt --encoder=./{sherpa}/{model}/encoder-epoch-99-avg-1.int8.onnx --decoder=./{sherpa}/{model}/decoder-epoch-99-avg-1.int8.onnx --joiner=./{sherpa}/{model}/joiner-epoch-99-avg-1.int8.onnx "
    num_utts = 0
    audios = []
    batch_size = 200
    complete = 0
    with open(args.scp_path, 'r', encoding='utf8') as scp, open(args.trans_path, 'w+', encoding='utf8') as trans:
        for line in [l.strip() for l in scp if l.strip()]:
            cols = line.split()
            if (len(cols) == 2):
                key, audio = cols
                print(F'{num_utts}\tkey:{key}\taudio:{audio}', file=sys.stderr, flush=True)
                #     text = recognize(audio)
                #     print(key + '\t' + text, file=trans, flush=True)
                audios.append(audio)
                num_utts += 1
            else:
                print(F'Invalid line: {line}', file=sys.stderr, flush=True)
        audio_slices = arr_size(audios, batch_size)
        for audio_slice in audio_slices:
            temp_cmd = cmd + ' '.join(audio_slice)
            complete += batch_size
            print(temp_cmd, file=sys.stderr, flush=True)
            print(F"progress: {complete}/{num_utts}", file=sys.stderr, flush=True)
            r = subprocess.run(temp_cmd, shell=True, capture_output=True, encoding='utf-8')
            print(r.stderr, file=sys.stderr, flush=True)
            lines = r.stderr.split("\n")
            for i in range(len(lines)):
                line = lines[i]
                if line.startswith("{"):
                    result = json.loads(line)
                    text = result['text']
                    key = lines[i - 3].split("/")[-1].split(".")[0]
                    print(key + '\t' + text, file=trans, flush=True)
